﻿#include "yolov7.h"
YOLOV7::YOLOV7(Net_config config, std::string yoloname) {
    this->confThreshold = config.confThreshold;
    this->nmsThreshold = config.nmsThreshold;

    std::string classesFile = "static/coco.names";
    std::string model_path = config.modelpath;
    std::wstring widestr = std::wstring(model_path.begin(), model_path.end());
    OrtStatus* status = OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0);
    sessionOptions.SetGraphOptimizationLevel(ORT_ENABLE_BASIC);
    ort_session = new Session(env, model_path.c_str(), sessionOptions);
    size_t numInputNodes = ort_session->GetInputCount();
    size_t numOutputNodes = ort_session->GetOutputCount();

    //---------------------------------替换GetInputName-----------------------------
    for (int i = 0; i < numInputNodes; i++) {
        AllocatorWithDefaultOptions allocator;
        In_AllocatedStringPtr.push_back(ort_session->GetInputNameAllocated(i, allocator));
        input_names.push_back(In_AllocatedStringPtr.at(i).get());
        Ort::TypeInfo input_type_info = ort_session->GetInputTypeInfo(i);
        auto input_tensor_info = input_type_info.GetTensorTypeAndShapeInfo();
        auto input_dims = input_tensor_info.GetShape();
        input_node_dims.push_back(input_dims);
    }
    for (int i = 0; i < numOutputNodes; i++) {
        AllocatorWithDefaultOptions allocator;
        Out_AllocatedStringPtr.push_back(ort_session->GetOutputNameAllocated(i, allocator));
        output_names.push_back(Out_AllocatedStringPtr.at(i).get());
        Ort::TypeInfo output_type_info = ort_session->GetOutputTypeInfo(i);
        auto output_tensor_info = output_type_info.GetTensorTypeAndShapeInfo();
        auto output_dims = output_tensor_info.GetShape();
        output_node_dims.push_back(output_dims);
    }
    //----------------------------------结束替换------------------------------------
    this->inpHeight = input_node_dims[0][2];
    this->inpWidth = input_node_dims[0][3];
    this->nout = output_node_dims[0][2];
    this->num_proposal = output_node_dims[0][1];

    std::ifstream ifs(classesFile.c_str());
    std::string line;
    while (getline(ifs, line))
        this->class_names.push_back(line);
    this->num_class = class_names.size();
}

void YOLOV7::normalize_(cv::Mat img) {
    //    img.convertTo(img, CV_32F);
    int row = img.rows;
    int col = img.cols;
    this->input_image_.resize(row * col * img.channels());
    for (int c = 0; c < 3; c++) {
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                float pix = img.ptr<uchar>(i)[j * 3 + 2 - c];
                this->input_image_[c * row * col + i * col + j] = pix / 255.0;
            }
        }
    }
}

void YOLOV7::nms(std::vector<BoxInfo>& input_boxes) {
    sort(input_boxes.begin(), input_boxes.end(), [](BoxInfo a, BoxInfo b) { return a.score > b.score; });
    std::vector<float> vArea(input_boxes.size());
    for (int i = 0; i < int(input_boxes.size()); ++i) {
        vArea[i] = (input_boxes.at(i).x2 - input_boxes.at(i).x1 + 1) * (input_boxes.at(i).y2 - input_boxes.at(i).y1 + 1);
    }

    std::vector<bool> isSuppressed(input_boxes.size(), false);
    for (int i = 0; i < int(input_boxes.size()); ++i) {
        if (isSuppressed[i]) {
            continue;
        }
        for (int j = i + 1; j < int(input_boxes.size()); ++j) {
            if (isSuppressed[j]) {
                continue;
            }
            float xx1 = (cv::max)(input_boxes[i].x1, input_boxes[j].x1);
            float yy1 = (cv::max)(input_boxes[i].y1, input_boxes[j].y1);
            float xx2 = (cv::min)(input_boxes[i].x2, input_boxes[j].x2);
            float yy2 = (cv::min)(input_boxes[i].y2, input_boxes[j].y2);

            float w = (cv::max)(float(0), xx2 - xx1 + 1);
            float h = (cv::max)(float(0), yy2 - yy1 + 1);
            float inter = w * h;
            float ovr = inter / (vArea[i] + vArea[j] - inter);

            if (ovr >= this->nmsThreshold) {
                isSuppressed[j] = true;
            }
        }
    }
    // return post_nms;
    int idx_t = 0;
    input_boxes.erase(remove_if(input_boxes.begin(), input_boxes.end(), [&idx_t, &isSuppressed](const BoxInfo& f) { return isSuppressed[idx_t++]; }),
                      input_boxes.end());
}

std::vector<BoxInfo> YOLOV7::detect(cv::Mat& frame) {
    cv::Mat dstimg;
    resize(frame, dstimg, cv::Size(this->inpWidth, this->inpHeight));
    this->normalize_(dstimg);
    std::array<std::int64_t, 4> input_shape_{1, 3, this->inpHeight, this->inpWidth};

    auto allocator_info = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Value input_tensor_ = Value::CreateTensor<float>(allocator_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());

    // 推理
    std::vector<Value> ort_outputs = ort_session->Run(RunOptions{nullptr}, &input_names[0], &input_tensor_, 1, output_names.data(), output_names.size());  // ��ʼ����
    /////generate proposals
    std::vector<BoxInfo> generate_boxes;
    float ratioh = (float)frame.rows / this->inpHeight, ratiow = (float)frame.cols / this->inpWidth;
    int n = 0, k = 0;  /// cx,cy,w,h,box_score, class_score
    const float* pdata = ort_outputs[0].GetTensorMutableData<float>();
    for (n = 0; n < this->num_proposal; n++)  /// ����ͼ�߶�
    {
        float box_score = pdata[4];
        if (box_score > this->confThreshold) {
            int max_ind = 0;
            float max_class_socre = 0;
            for (k = 0; k < num_class; k++) {
                if (pdata[k + 5] > max_class_socre) {
                    max_class_socre = pdata[k + 5];
                    max_ind = k;
                }
            }
            max_class_socre *= box_score;
            if (max_class_socre > this->confThreshold) {
                float cx = pdata[0] * ratiow;  /// cx
                float cy = pdata[1] * ratioh;  /// cy
                float w = pdata[2] * ratiow;   /// w
                float h = pdata[3] * ratioh;   /// h

                float xmin = cx - 0.5 * w;
                float ymin = cy - 0.5 * h;
                float xmax = cx + 0.5 * w;
                float ymax = cy + 0.5 * h;

                generate_boxes.push_back(BoxInfo{xmin, ymin, xmax, ymax, max_class_socre, max_ind});
            }
        }
        pdata += nout;
    }

    // Perform non maximum suppression to eliminate redundant overlapping boxes with
    // lower confidences
    nms(generate_boxes);

    /*for (size_t i = 0; i < generate_boxes.size(); ++i)
    {
            int xmin = int(generate_boxes[i].x1);
            int ymin = int(generate_boxes[i].y1);
            cv::rectangle(frame, cv::Point(xmin, ymin), cv::Point(int(generate_boxes[i].x2), int(generate_boxes[i].y2)), cv::Scalar(0, 0, 255), 2);
            std::string label = cv::format("%.2f", generate_boxes[i].score);
            label = this->class_names[generate_boxes[i].label] + ":" + label;
            cv::putText(frame, label, cv::Point(xmin, ymin - 5), cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0, 255, 0), 1);
    }*/
    return generate_boxes;
}
